/////////////////////////////////////////////////////////////////////////////
// Original code from libhdfs3. Copyright (c) 2013 - 2014, Pivotal Inc.
// All rights reserved. Author: Zhanwei Wang
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <ktest/ktest.h>
#include <kmock/kmock.h>

#include <kmhdfs/client/input_stream_impl.h>
#include <kmhdfs/client/file_system.h>
#include <tests/mock/mock_file_system_inter.h>
#include <tests/mock/test_datanode_stub.h>
#include <tests/mock/mock_datanode.h>
#include <kmhdfs/server/extended_block.h>
#include <kmhdfs/common/xml_config.h>
#include <string>

using namespace Hdfs::Internal;
using namespace Hdfs::Mock;
using namespace testing;
using ::testing::AtLeast;

class TestInputSteam: public ::testing::Test {

public:
    TestInputSteam() {
    }
    ~TestInputSteam() {
    }

protected:
    MockFileSystemInter fs;
    InputStreamImpl ins;
};

class MockLocatedBlocks: public LocatedBlocks {
public:
    MOCK_CONST_METHOD0(getFileLength, int64_t ());
    MOCK_METHOD1(setFileLength, void (int64_t fileLength));
    MOCK_CONST_METHOD0(isLastBlockComplete, bool ());
    MOCK_METHOD1(setIsLastBlockComplete, void (bool lastBlockComplete));
    MOCK_METHOD0(getLastBlock, std::shared_ptr<LocatedBlock> ());
    MOCK_METHOD1(setLastBlock, void (const LocatedBlock & lastBlock));
    MOCK_CONST_METHOD0(isUnderConstruction, bool ());
    MOCK_METHOD1(setUnderConstruction, void (bool underConstruction));
    MOCK_METHOD1(findBlock, LocatedBlock * (int64_t position));
    MOCK_METHOD2(findBlock, LocatedBlock * (int64_t position, int32_t & targetBlockIdx));
    MOCK_METHOD0(getBlocks, std::vector<LocatedBlock> & ());
    MOCK_METHOD1(setLastBlock, void(std::shared_ptr<LocatedBlock>));
    MOCK_CONST_METHOD0(getEcPolicy, std::shared_ptr<ECPolicy> ());
    MOCK_METHOD1(setEcPolicy, void (std::shared_ptr<ECPolicy>));
    MOCK_METHOD2(insertRange, void (int32_t blockIdx, std::vector<LocatedBlock> & newBlocks));
    MOCK_METHOD5(addAll, void (std::vector<LocatedBlock> & oldBlocks, int32_t index, std::vector<LocatedBlock> & newBlocks,
            int32_t start, int32_t end));
};

class MockDatanodeStub: public TestDatanodeStub {
public:
    MOCK_METHOD0(getDatanode, std::shared_ptr<MockDatanode>());

};

TEST(InputStreamTest, ReadBlockLength_Success) {
    InputStreamImpl ins;
    LocatedBlock b;
    std::shared_ptr<MockDatanode> datanode(new MockDatanode());
    MockDatanodeStub stub;
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df;
    dfv.push_back(df);
    b.locs = dfv.get();
    ins.stub = &stub;
    EXPECT_CALL(stub, getDatanode()).Times(1).WillOnce(Return(datanode));
    EXPECT_CALL(*datanode, getReplicaVisibleLength(_)).Times(1).WillOnce(Return(2));
    EXPECT_NO_THROW(ins.readBlockLength(b));
}

TEST(InputStreamTest, ReadBlockLength_Fail) {
    InputStreamImpl ins;
    std::shared_ptr<MockDatanode> datanode(new MockDatanode());
    MockDatanodeStub stub;
    LocatedBlock b;
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df1;
    dfv.push_back(df1);
    b.locs = dfv.get();
    ins.stub = &stub;
    EXPECT_CALL(stub, getDatanode()).Times(1).WillOnce(Return(datanode));
    EXPECT_CALL(*datanode, getReplicaVisibleLength(_)).Times(1).WillOnce(Return(-1));
    EXPECT_EQ(ins.readBlockLength(b), -1);
}

TEST(InputStreamTest, ReadBlockLength_FailTry) {
    InputStreamImpl ins;
    std::shared_ptr<MockDatanode> datanode(new MockDatanode());
    MockDatanodeStub stub;
    LocatedBlock b;
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df1;
    DatanodeInfo df2;
    dfv.push_back(df1);
    dfv.push_back(df2);
    b.locs = dfv.get();
    ins.stub = &stub;
    EXPECT_CALL(stub, getDatanode()).Times(2).WillOnce(Return(datanode)).WillOnce(Return(datanode));
    EXPECT_CALL(*datanode, getReplicaVisibleLength(_)).Times(2).WillOnce(Return(-1)).WillOnce(
        Return(0));
    EXPECT_EQ(ins.readBlockLength(b), 0);
}

TEST(InputStreamTest, ReadBlockLength_ThrowTry) {
    InputStreamImpl ins;
    std::shared_ptr<MockDatanode> datanode(new MockDatanode());
    MockDatanodeStub stub;
    LocatedBlock b;
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df1;
    DatanodeInfo df2;
    dfv.push_back(df1);
    dfv.push_back(df2);
    b.locs = dfv.get();
    ins.stub = &stub;
    Hdfs::ReplicaNotFoundException e("test", "test", 2, "test");
    EXPECT_CALL(stub, getDatanode()).Times(2).WillOnce(Return(datanode)).WillOnce(Return(datanode));
    EXPECT_CALL(*datanode, getReplicaVisibleLength(_)).Times(2).WillOnce(Throw(e)).WillOnce(
        Return(0));
    EXPECT_EQ(ins.readBlockLength(b), 0);
}

TEST(InputStreamTest, UpdateBlockInfos_LastComplete) {
    MockFileSystemInter * fs = new MockFileSystemInter;
    MockLocatedBlocks * lbs = new MockLocatedBlocks();
    InputStreamImpl ins;
    ins.filesystem = std::shared_ptr<FileSystemInter>(fs);
    ins.maxGetBlockInfoRetry = 1;
    ins.lastBlockBeingWrittenLength = 1;
    ins.lbs = std::shared_ptr < MockLocatedBlocks > (lbs);
    EXPECT_CALL(*fs, getBlockLocations(_, _, _, _)).WillOnce(Return());
    EXPECT_CALL(*lbs, isLastBlockComplete()).WillOnce(Return(true));
    EXPECT_NO_THROW(ins.updateBlockInfos());
    EXPECT_EQ(ins.lastBlockBeingWrittenLength, 0);
}

TEST(InputStreamTest, UpdateBlockInfos_NotLastComplete) {
    MockFileSystemInter * fs = new MockFileSystemInter;
    MockLocatedBlocks * lbs = new MockLocatedBlocks();
    InputStreamImpl ins;
    ins.filesystem = std::shared_ptr<FileSystemInter>(fs);
    ins.maxGetBlockInfoRetry = 1;
    ins.lastBlockBeingWrittenLength = 1;
    ins.lbs = std::shared_ptr < MockLocatedBlocks > (lbs);
    std::shared_ptr<MockDatanode> datanode(new MockDatanode());
    MockDatanodeStub stub;
    ins.stub = &stub;
    std::shared_ptr<LocatedBlock> block(new LocatedBlock);
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df;
    dfv.push_back(df);
    block->setLocations(dfv.get());
    EXPECT_CALL(*fs, getBlockLocations(_, _, _, _)).WillOnce(Return());
    EXPECT_CALL(*lbs, isLastBlockComplete()).WillOnce(Return(false));
    EXPECT_CALL(*lbs, getLastBlock()).WillOnce(Return(block));
    EXPECT_CALL(stub, getDatanode()).Times(1).WillOnce(Return(datanode));
    EXPECT_CALL(*datanode, getReplicaVisibleLength(_)).WillOnce(Return(2));
    EXPECT_NO_THROW(ins.updateBlockInfos());
    EXPECT_GT(ins.lastBlockBeingWrittenLength, 0);
}

TEST(InputStreamTest, UpdateBlockInfos_TrowAndTry) {
    MockFileSystemInter * fs = new MockFileSystemInter;
    MockLocatedBlocks * lbs = new MockLocatedBlocks();
    InputStreamImpl ins;
    ins.filesystem = std::shared_ptr<FileSystemInter>(fs);
    ins.maxGetBlockInfoRetry = 2;
    ins.lastBlockBeingWrittenLength = 1;
    ins.lbs = std::shared_ptr < MockLocatedBlocks > (lbs);
    Hdfs::HdfsRpcException e("test", "test", 2, "test");
    EXPECT_CALL(*fs, getBlockLocations(_, _, _, _)).Times(2).WillOnce(Throw(e)).WillOnce(
        Throw(e));
    EXPECT_THROW(ins.updateBlockInfos(), Hdfs::HdfsRpcException);
}

TEST(InputStreamTest, ChoseBestNode_Success) {
    InputStreamImpl ins;
    LocatedBlock * lb = new LocatedBlock();
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df1;
    DatanodeInfo df2;
    dfv.push_back(df1);
    dfv.push_back(df2);
    lb->locs = dfv.get();
    ins.curBlock = std::shared_ptr < LocatedBlock > (lb);
    EXPECT_NO_THROW(ins.choseBestNode());
}

TEST(InputStreamTest, SetupBlockReader_Failed) {
    InputStreamImpl ins;
    LocatedBlock * lb = new LocatedBlock();
    LockVector<DatanodeInfo> dfv;
    DatanodeInfo df1;
    DatanodeInfo df2;
    dfv.push_back(df1);
    dfv.push_back(df2);
    lb->locs = dfv.get();
    ins.curBlock = std::shared_ptr < LocatedBlock > (lb);
    std::vector<DatanodeInfo> tmp = dfv.get();
    ins.failedNodes.set(tmp);
    EXPECT_THROW(ins.setupBlockReader(false), Hdfs::HdfsIOException);
}
